'''
Created on 21 juil. 2010

@author: laurent
'''
import commands, threading, fcntl, signal, os
import struct, time, termios, pty, sys, re, select

from SynchronizedMethod import SynchronizedMethod
from Terminal import Terminal

class Multiplex:
    def __init__(self, cmd = None, env_term = None):
        # Set Linux signal handler
        uname = commands.getoutput('uname')
        if uname == 'Linux':
            self.sigchldhandler = signal.signal(signal.SIGCHLD, signal.SIG_IGN)
        # Session
        self.session = {}
        self.cmd = cmd
        self.env_term = env_term
        # Synchronize methods
        self.lock = threading.RLock()
        for name in ['proc_keepalive', 'proc_buryall',
            'proc_read', 'proc_write', 'proc_dump', 'proc_getalive']:
            orig = getattr(self, name)
            setattr(self, name, SynchronizedMethod(self.lock, orig))
        # Supervisor thread
        self.signal_stop = 0
        self.thread = threading.Thread(target = self.proc_thread)
        self.thread.start()
    def stop(self):
        # Stop supervisor thread
        self.signal_stop = 1
        self.thread.join()
    def proc_keepalive(self, sid, w, h):
        if not sid in self.session:
            # Start a new session
            self.session[sid] = {
                'state':'unborn',
                'term':    Terminal(w, h),
                'time':    time.time(),
                'w':    w,
                'h':    h}
            return self.proc_spawn(sid)
        elif self.session[sid]['state'] == 'alive':
            self.session[sid]['time'] = time.time()
            # Update terminal size
            if self.session[sid]['w'] != w or self.session[sid]['h'] != h:
                try:
                    fcntl.ioctl(self.session[sid]['fd'], 
                        struct.unpack('i', 
                            struct.pack('I', termios.TIOCSWINSZ)
                        )[0],
                        struct.pack("HHHH", h, w, 0, 0))
                except (IOError, OSError):
                    pass
                self.session[sid]['term'].set_size(w, h)
                self.session[sid]['w'] = w
                self.session[sid]['h'] = h
            return True
        else:
            return False
    def proc_spawn(self, sid):
        # Session
        self.session[sid]['state'] = 'alive'
        w, h = self.session[sid]['w'], self.session[sid]['h']
        # Fork new process
        try:
            pid, fd = pty.fork()
        except (IOError, OSError):
            self.session[sid]['state'] = 'dead'
            return False
        if pid == 0:
            if self.cmd:
                cmd = self.cmd
            else:
                sys.stdout.write("Login: ")
                login = sys.stdin.readline().strip()
                if re.match('^[0-9A-Za-z-_.]+$', login):
                    cmd = 'ssh'
                    cmd += ' -oPreferredAuthentications=keyboard-interactive,password'
                    cmd += ' -oNoHostAuthenticationForLocalhost=yes'
                    cmd += ' -oLogLevel=FATAL'
                    cmd += ' -F/dev/null -l' + login +' localhost'
                else:
                    os._exit(0)
            # Safe way to make it work under BSD and Linux
            try:
                ls = os.environ['LANG'].split('.')
            except KeyError:
                ls = []
            if len(ls) < 2:
                ls = ['en_US', 'UTF-8']
            try:
                os.putenv('COLUMNS', str(w))
                os.putenv('LINES', str(h))
                os.putenv('TERM', self.env_term)
                os.putenv('PATH', os.environ['PATH'])
                os.putenv('LANG', ls[0] + '.UTF-8')
                os.system(cmd)
            except (IOError, OSError):
                pass
#            self.proc_finish(sid)
            os._exit(0)
        else:
            # Store session vars
            self.session[sid]['pid'] = pid
            self.session[sid]['fd'] = fd
            # Set file control
            fcntl.fcntl(fd, fcntl.F_SETFL, os.O_NONBLOCK)
            # Set terminal size
            try:
                fcntl.ioctl(fd, 
                    struct.unpack('i', 
                        struct.pack('I', termios.TIOCSWINSZ)
                    )[0],
                    struct.pack("HHHH", h, w, 0, 0))
            except (IOError, OSError):
                pass
            return True
    def proc_waitfordeath(self, sid):
        try:
            os.close(self.session[sid]['fd'])
        except (KeyError, IOError, OSError):
            pass
        if sid in self.session:
            if 'fd' in self.session[sid]:
                del self.session[sid]['fd']
        try:
            os.waitpid(self.session[sid]['pid'], 0)
        except (KeyError, IOError, OSError):
            pass
        if sid in self.session:
            if 'pid' in self.session[sid]:
                del self.session[sid]['pid']
        self.session[sid]['state'] = 'dead'
        return True
    def proc_bury(self, sid):
        if self.session[sid]['state'] == 'alive':
            try:
                os.kill(self.session[sid]['pid'], signal.SIGTERM)
            except (IOError, OSError):
                pass
        self.proc_waitfordeath(sid)
        if sid in self.session:
            del self.session[sid]
        return True
    def proc_buryall(self):
        for sid in self.session.keys():
            self.proc_bury(sid)
            
    # Read from process
    def proc_read(self, sid):
        if sid not in self.session:
            return False
        elif self.session[sid]['state'] != 'alive':
            return False
        try:
            fd = self.session[sid]['fd']
            d = os.read(fd, 65536)
            if not d:
                # Process finished, BSD
                self.proc_waitfordeath(sid)
                return False
        except (IOError, OSError):
            # Process finished, Linux
            self.proc_waitfordeath(sid)
            return False
        term = self.session[sid]['term']
        term.write(d)
        # Read terminal response
        d = term.read()
        if d:
            try:
                os.write(fd, d)
            except (IOError, OSError):
                return False
        return True
    # Write to process
    def proc_write(self, sid, d):
        if sid not in self.session:
            return False
        elif self.session[sid]['state'] != 'alive':
            return False
        try:
            term = self.session[sid]['term']
            d = term.pipe(d)
            fd = self.session[sid]['fd']
            os.write(fd, d)
        except (IOError, OSError):
            return False
        return True
    # Dump terminal output
    def proc_dump(self, sid):
        if sid not in self.session:
            return False
        return self.session[sid]['term'].dump()
        
    # Get alive sessions, bury timed out ones
    def proc_getalive(self):
        fds = []
        fd2sid = {}
        now = time.time()
        for sid in self.session.keys():
            then = self.session[sid]['time']
            if (now - then) > 60:
                self.proc_bury(sid)
            else:
                if self.session[sid]['state'] == 'alive':
                    fds.append(self.session[sid]['fd'])
                    fd2sid[self.session[sid]['fd']] = sid
        return (fds, fd2sid)
    # Supervisor thread
    def proc_thread(self):
        while not self.signal_stop:
            # Read fds
            (fds, fd2sid) = self.proc_getalive()
            try:
                i, o, e = select.select(fds, [], [], 1.0)
            except (IOError, OSError):
                i = []
            for fd in i:
                sid = fd2sid[fd]
                self.proc_read(sid)
            if len(i):
                time.sleep(0.002)
        self.proc_buryall()